import os

from keras.saving.saved_model.load import metrics
from tensorflow import saved_model

from prepare_mnist_data import load_mnist_data, process_mnist_data

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEl'] = '2'
    (train_images, train_labels), (test_images, test_labels) = load_mnist_data()

    db, ds_val = process_mnist_data(train_images, train_labels, test_images, test_labels, 128)

    model = saved_model.load("mnist_savemodel")
    acc_meter = metrics.CategoricalAccuracy()

    for x, y in ds_val:
        pred = model(x)
        acc_meter.update_state(y_true=y, y_pred=pred)

    print("Test Accuracy: %f" % acc_meter.result())
